#-*-coding:utf-8-*-#

import os
import pandas as pd
import matplotlib.pyplot as plt

currentPath = os.getcwd()
os.chdir('PATH')

df = pd.read_csv('Data_Fig1a_1.csv') # US Drought Monitor data

TW = pd.read_csv('Data_Fig1a_2.csv') # Drought Twitter data with negative sentiment and network information
dates = TW['date']
dates = pd.to_datetime(dates, errors='coerce')
dates = dates.dt.strftime("%Y-%m-%d")


df0 = pd.DataFrame()
df0['Date'] = df.ValidEnd
df0['D0D4'] = df.D0 + df.D1 + df.D2 + df.D3 + df.D4
df0['D1D4'] = df.D1 + df.D2 + df.D3 + df.D4
df0['D2D4'] = df.D2 + df.D3 + df.D4
df0['D3D4'] = df.D3 + df.D4
df0['D4'] = df.D4

df0['Date'] = pd.to_datetime(df0['Date'], errors='coerce')
df0['Date'] = df0['Date'].dt.strftime("%Y-%m-%d")
df0 = df0.sort_values(by='Date')
df0 = df0.reset_index(drop=True)

df['ValidEnd'] = pd.to_datetime(df['ValidEnd'], errors='coerce')
df['ValidEnd'] = df['ValidEnd'].dt.strftime("%Y-%m-%d")
df = df.sort_values(by='ValidEnd')
df = df.reset_index(drop=True)


"""

Transform to weekly data

"""
n_TW = pd.DataFrame()
z=0
for i in range(2,len(TW),7):
    n_TW.loc[z,'date'] = TW['date'].iloc[i+3]
    n_TW.loc[z,'count'] = TW[i:i+7]['count'].sum()
    z +=1


fig, ax1 = plt.subplots(dpi=400,figsize=(6,2))
labels = ['D0','D1','D2','D3','D4']
pal = ['#730000','#E60000','#FFAA00','#FCD37F','#FFFF00']
ax1.stackplot(df0.Date[9:],df.D4[9:],df.D3[9:],df.D2[9:],df.D1[9:],df.D0[9:],colors = pal,zorder=2 ,alpha=0.7)
ax1.set_xticks(range(0,98,10))
ax1.set_xlim([0,len(n_TW)-1])
ax1.set_ylim([0,100])
ax1.set_xlabel('Date')
ax1.axhline(y=20,color='black',linestyle=':',zorder=1)
ax1.axhline(y=40,color='black',linestyle=':',zorder=1)
ax1.axhline(y=60,color='black',linestyle=':',zorder=1)
ax1.axhline(y=80,color='black',linestyle=':',zorder=1)

ax2 = ax1.twinx()
ax2.set_ylabel('Twitter',color='blue',rotation=270,va='bottom')
line2 = ax2.plot(n_TW['date'],n_TW['count'],color='blue',linewidth=1.0)
ax2.tick_params(axis='y', labelcolor='blue')
plt.axvline(x=18,linestyle='--',color='k',linewidth=1.0)
plt.text(13,2500,'July 10')
plt.tight_layout()
plt.show()


